# ==========================================================================
# Eses Composition Golden Ratio
# ==========================================================================
#
# Description:
# The 'Eses Composition Golden Ratio' is a non-destructive visualization
# tool for ComfyUI. It overlays a highly customizable golden ratio pattern
# (Fibonacci spiral) onto a preview image, serving as a powerful guide for
# analyzing and creating harmoniously balanced compositions. The overlay does
# not alter the final output image, making it a pure guidance utility.
#
# Key Features:
#
# - Pattern Generation:
#   - Orientation: Sets the starting direction of the pattern (Up, Right,
#     Down, Left) with an 'auto' mode that adapts to image dimensions.
#   - Steps: Controls the number of recursive divisions in the pattern.
#   - Draw Spiral: Toggles the visibility of the spiral curve itself.
#
# - Fitting & Sizing:
#   - Fit Mode: 'Crop' maintains the perfect golden ratio, potentially
#     leaving empty space, while 'Stretch' fits the pattern to the
#     image's aspect ratio.
#   - Crop Offset: When in 'Crop' mode, adjusts the pattern's position
#     within the image frame.
#   - Axial Stretch: Manually stretches or squashes the pattern along its
#     main axis for artistic adjustments.
#
# - Projection & Transforms:
#   - Offset X/Y: Freely moves the entire pattern across the image.
#   - Rotation: Rotates the pattern from -360 to 360 degrees.
#   - Scale: Uniformly zooms the pattern in or out from its center.
#   - Flip Horizontal/Vertical: Flips the pattern's orientation.
#
# - Line & Style Settings:
#   - Line Color: Sets the color of the pattern overlay.
#   - Line Thickness: Controls the base thickness of the lines.
#   - Uniform Line Width: When enabled, prevents line thickness from
#     changing during uniform scaling.
#   - Blend Mode: Sets the canvas blend mode for the overlay effect.
#
# Usage:
# Connect an image to the 'image' input. The guide will appear as an
# overlay on the preview image within the node itself. Adjust the various
#
# parameters to manipulate the guide's position, size, and appearance in
# real-time. The output 'image' tensor remains unaltered.
#
# Version: 1.0.2
#
# License: See LICENSE.txt
#
# ==========================================================================

import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
from server import PromptServer # type: ignore
from io import BytesIO
import base64

class EsesCompositionGoldenRatio:
    """
    A non-destructive visualization tool for ComfyUI. It overlays a golden
    ratio compositional guide onto a preview image without altering the output.
    """

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    @classmethod
    def INPUT_TYPES(cls):
        """
        Defines the input types for the node for the Golden Ratio guide.
        """
        
        blend_modes = [
            "source-over", "lighter", "screen", "multiply", "overlay", "darken",
            "lighten", "color-dodge", "color-burn", "hard-light", "soft-light",
            "difference", "exclusion", "hue", "saturation", "color", "luminosity"
        ]
        
        return {
            "required": {
                "image": ("IMAGE",),

                # --- General Settings ---
                "preview_resolution_limit": ("INT", {"default": 1024, "min": 256, "max": 8192, "step": 64}),
                "line_color_rgb": ("STRING", {"default": "255,255,255,255"}), # Renamed for clarity
                "line_thickness": ("FLOAT", {"default": 1.5, "min": 0.1, "max": 32.0, "step": 0.1, "round": 0.01}),
                "uniform_line_width": ("BOOLEAN", {"default": False, "label_on": "on", "label_off": "off"}), # <-- ADD THIS
                "blend_mode": (blend_modes,),
                
                # --- Golden Ratio Settings ---
                "orientation": (["auto", "up", "down", "left", "right"],),
                "fit_mode": (["crop", "stretch"], {"default": "crop"}),
                "crop_offset": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01}), # <-- ADD THIS LINE
                "flip_horizontal": ("BOOLEAN", {"default": False, "label_on": "on", "label_off": "off"}),
                "flip_vertical": ("BOOLEAN", {"default": False, "label_on": "on", "label_off": "off"}),
                "draw_spiral": ("BOOLEAN", {"default": True, "label_on": "on", "label_off": "off"}),
                "steps": ("INT", {"default": 8, "min": 1, "max": 20, "step": 1}),
                "axial_stretch": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 4.0, "step": 0.01, "round": 0.01}), # <-- ADD THIS
                
                # --- Projection Transforms ---
                "offset_x": ("FLOAT", {"default": 0.0, "min": -4096, "max": 4096, "step": 1.0}),
                "offset_y": ("FLOAT", {"default": 0.0, "min": -4096, "max": 4096, "step": 1.0}),
                "rotation": ("FLOAT", {"default": 0.0, "min": -360.0, "max": 360.0, "step": 0.1, "round": 0.01}),
                "scale": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 5.0, "step": 0.01, "round": 0.01}),
            },
            "optional": {
                "mask": ("MASK",),
            },
            "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
        }
    

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "mask")
    FUNCTION = "execute"
    CATEGORY = "Eses Nodes/Visualization"

    def execute(self, image: torch.Tensor, preview_resolution_limit, line_color_rgb, line_thickness, uniform_line_width, blend_mode,
            orientation, fit_mode, crop_offset, flip_horizontal, flip_vertical, draw_spiral, steps, axial_stretch,
            offset_x, offset_y, rotation, scale,
            mask=None, prompt=None, extra_pnginfo=None):
        
        image_for_preview = image
        
        _, h, w, _ = image_for_preview.shape
        
        if max(h, w) > preview_resolution_limit:
            if w > h:
                new_w = preview_resolution_limit
                new_h = int(h * (preview_resolution_limit / w))
            else:
                new_h = preview_resolution_limit
                new_w = int(w * (preview_resolution_limit / h))
                
            image_for_preview_permuted = image_for_preview.permute(0, 3, 1, 2)
            resized_preview = F.interpolate(image_for_preview_permuted, size=(new_h, new_w), mode='bilinear', align_corners=False)
            image_for_preview = resized_preview.permute(0, 2, 3, 1)

        preview_image_tensor = image_for_preview[0]
        i = 255. * preview_image_tensor.cpu().numpy()
        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))

        buffered = BytesIO()
        img.save(buffered, format="PNG")
        img_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
        node_id = None

        if prompt:
            for k, v in prompt.items():
                if v.get('class_type') == type(self).__name__:
                    node_id = k
                    break
        
        if node_id:
            PromptServer.instance.send_sync("eses.composition_golden_ratio_preview", {
                "node_id": node_id,
                "image_data": img_base64,
            })

        return (image, mask)
